import os
import time
import requests
import logging
from typing import Dict, Optional, List
from enum import Enum

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class InstanceStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    STOPPED = "stopped"
    ERROR = "error"

class AutoDLClient:
    """Client for interacting with AutoDL API for elastic deployment"""
    
    def __init__(self, api_key: str, endpoint: str = "https://api.autodl.com"):
        """Initialize the AutoDL client"""
        self.api_key = api_key
        self.endpoint = endpoint.rstrip('/')
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        })
    
    def create_instance(self, config: Dict) -> Dict:
        """Create a new instance with the given configuration"""
        url = f"{self.endpoint}/v1/instances"
        try:
            response = self.session.post(url, json=config)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            logger.error(f"Error creating instance: {e}")
            raise
    
    def get_instance(self, instance_id: str) -> Dict:
        """Get details of a specific instance"""
        url = f"{self.endpoint}/v1/instances/{instance_id}"
        try:
            response = self.session.get(url)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            logger.error(f"Error getting instance {instance_id}: {e}")
            raise
    
    def list_instances(self, status: Optional[str] = None) -> List[Dict]:
        """List all instances, optionally filtered by status"""
        url = f"{self.endpoint}/v1/instances"
        params = {}
        if status:
            params["status"] = status
        
        try:
            response = self.session.get(url, params=params)
            response.raise_for_status()
            return response.json().get("instances", [])
        except requests.exceptions.RequestException as e:
            logger.error(f"Error listing instances: {e}")
            raise
    
    def stop_instance(self, instance_id: str) -> bool:
        """Stop a running instance"""
        url = f"{self.endpoint}/v1/instances/{instance_id}/stop"
        try:
            response = self.session.post(url)
            response.raise_for_status()
            return True
        except requests.exceptions.RequestException as e:
            logger.error(f"Error stopping instance {instance_id}: {e}")
            return False
    
    def terminate_instance(self, instance_id: str) -> bool:
        """Terminate an instance"""
        url = f"{self.endpoint}/v1/instances/{instance_id}"
        try:
            response = self.session.delete(url)
            response.raise_for_status()
            return True
        except requests.exceptions.RequestException as e:
            logger.error(f"Error terminating instance {instance_id}: {e}")
            return False
    
    def wait_for_instance_status(
        self, 
        instance_id: str, 
        target_status: InstanceStatus, 
        timeout: int = 300,
        check_interval: int = 5
    ) -> bool:
        """Wait for an instance to reach a specific status"""
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            try:
                instance = self.get_instance(instance_id)
                current_status = instance.get("status")
                
                if current_status == target_status:
                    return True
                elif current_status == InstanceStatus.ERROR:
                    logger.error(f"Instance {instance_id} entered error state")
                    return False
                    
                time.sleep(check_interval)
            except Exception as e:
                logger.warning(f"Error checking instance status: {e}")
                time.sleep(check_interval)
        
        logger.warning(f"Timeout waiting for instance {instance_id} to reach status {target_status}")
        return False

class WanInstanceManager:
    """Manager for Wan 2.2 instances on AutoDL"""
    
    def __init__(self, autodl_client: AutoDLClient, max_instances: int = 2):
        self.client = autodl_client
        self.max_instances = max_instances
        self.wan_image = "wan-2.2"
    
    def get_active_instances(self) -> List[Dict]:
        """Get all active Wan instances"""
        instances = self.client.list_instances(status=InstanceStatus.RUNNING)
        return [i for i in instances if i.get("image") == self.wan_image]
    
    def can_scale_up(self) -> bool:
        """Check if we can start more instances"""
        active_instances = self.get_active_instances()
        return len(active_instances) < self.max_instances
    
    def start_instance(self, config_override: Dict = None) -> Optional[Dict]:
        """Start a new Wan instance"""
        if not self.can_scale_up():
            logger.warning("Maximum number of instances already running")
            return None
        
        default_config = {
            "image": self.wan_image,
            "instance_type": "GPU_RTX_3090",  # Adjust based on availability
            "name": f"wan-instance-{int(time.time())}",
            "auto_stop": True,
            "auto_stop_time": 3600,  # Auto stop after 1 hour of inactivity
            "env_vars": {
                "API_KEY": os.getenv("WAN_API_KEY", ""),
                "MODE": "api"
            }
        }
        
        # Apply any overrides
        config = {**default_config, **(config_override or {})}
        
        try:
            instance = self.client.create_instance(config)
            logger.info(f"Started new instance: {instance['id']}")
            
            # Wait for instance to be running
            if self.client.wait_for_instance_status(instance['id'], InstanceStatus.RUNNING):
                return instance
            return None
            
        except Exception as e:
            logger.error(f"Failed to start instance: {e}")
            return None
    
    def stop_idle_instances(self, max_idle_time: int = 600) -> int:
        """Stop instances that have been idle for too long"""
        stopped_count = 0
        instances = self.get_active_instances()
        
        for instance in instances:
            last_activity = instance.get('last_activity', 0)
            idle_time = time.time() - last_activity
            
            if idle_time > max_idle_time:
                if self.client.stop_instance(instance['id']):
                    stopped_count += 1
                    logger.info(f"Stopped idle instance: {instance['id']}")
        
        return stopped_count
